"""
数据库迁移脚本
将现有的JSON文件数据迁移到Supabase数据库

使用方法:
python migrate_to_database.py [data_dir] [--force]

参数:
- data_dir: 数据目录路径，默认为 "./data"  
- --force: 强制覆盖已存在的数据库数据
"""

import asyncio
import argparse
import os
import json
import sys
from typing import Dict, List, Optional
from datetime import datetime
from pathlib import Path

# 添加当前目录到路径，以便导入模块
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from database import get_db_service, RoleTrait
from database.supabase_client import check_supabase_config, test_connection

# 默认用户ID（单用户模式）
DEFAULT_USER_ID = '00000000-0000-0000-0000-000000000000'

class DataMigrator:
    """数据迁移器"""
    
    def __init__(self, data_dir: str, force: bool = False):
        self.data_dir = Path(data_dir)
        self.force = force
        self.stats = {
            'sessions_migrated': 0,
            'messages_migrated': 0,
            'roles_migrated': 0,
            'sessions_skipped': 0,
            'roles_skipped': 0,
            'errors': []
        }
    
    async def migrate(self):
        """执行完整迁移"""
        print("=== MindMate 数据库迁移工具 ===\n")
        
        # 1. 检查配置和连接
        if not await self._check_prerequisites():
            return False
            
        # 2. 检查数据目录
        if not self._check_data_directory():
            return False
            
        # 3. 执行迁移
        print("开始数据迁移...\n")
        
        try:
            # 迁移自定义角色
            await self._migrate_custom_roles()
            
            # 迁移会话和消息
            await self._migrate_sessions()
            
            # 显示迁移结果
            self._show_migration_results()
            
            return True
            
        except Exception as e:
            print(f"❌ 迁移过程中发生错误: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    async def _check_prerequisites(self) -> bool:
        """检查前置条件"""
        print("1. 检查Supabase配置...")
        if not check_supabase_config():
            print("❌ Supabase配置不完整")
            print("\n请按以下步骤配置：")
            print("1. 创建Supabase项目: https://supabase.com")
            print("2. 在项目设置中获取URL和anon key")
            print("3. 设置环境变量:")
            print("   set SUPABASE_URL=你的项目URL")
            print("   set SUPABASE_ANON_KEY=你的匿名密钥")
            print("4. 在SQL编辑器中执行 backend/database/init.sql")
            return False
        print("✅ 配置检查通过")
        
        print("2. 测试数据库连接...")
        if not test_connection():
            print("❌ 数据库连接失败")
            return False
        print("✅ 数据库连接成功")
        
        print("3. 初始化数据库服务...")
        db_service = get_db_service()
        if not db_service.initialize():
            print("❌ 数据库服务初始化失败")
            return False
        print("✅ 数据库服务初始化成功\n")
        
        # 存储服务实例供后续使用
        self.db_service = db_service
        
        return True
    
    def _check_data_directory(self) -> bool:
        """检查数据目录"""
        if not self.data_dir.exists():
            print(f"❌ 数据目录不存在: {self.data_dir}")
            return False
            
        print(f"数据目录: {self.data_dir}")
        
        # 检查可用文件
        sessions_file = self.data_dir / "sessions.json"
        roles_file = self.data_dir / "custom_roles.json"
        
        if sessions_file.exists():
            print(f"[OK] 找到会话文件: {sessions_file}")
        else:
            print(f"[WARN] 未找到会话文件: {sessions_file}")
            
        if roles_file.exists():
            print(f"[OK] 找到角色文件: {roles_file}")
        else:
            print(f"[WARN] 未找到角色文件: {roles_file}")
            
        if not sessions_file.exists() and not roles_file.exists():
            print("❌ 未找到任何可迁移的数据文件")
            return False
            
        print()
        return True
    
    async def _migrate_custom_roles(self):
        """迁移自定义角色"""
        print("迁移自定义角色...")
        
        roles_file = self.data_dir / "custom_roles.json"
        if not roles_file.exists():
            print("⚠️  跳过角色迁移 - 文件不存在")
            return
            
        try:
            with open(roles_file, 'r', encoding='utf-8') as f:
                roles_data = json.load(f)
                
            for role_id, role_data in roles_data.items():
                try:
                    # 检查角色是否已存在
                    if not self.force:
                        existing_role = await self.db_service.get_custom_role(role_id)
                        if existing_role:
                            print(f"  ⏭️  跳过角色 '{role_data.get('name', role_id)}' - 已存在")
                            self.stats['roles_skipped'] += 1
                            continue
                    
                    # 转换角色特征
                    traits = []
                    if 'custom_traits' in role_data and role_data['custom_traits']:
                        for trait in role_data['custom_traits']:
                            if isinstance(trait, dict) and 'key' in trait and 'value' in trait:
                                traits.append(RoleTrait(key=trait['key'], value=trait['value']))
                    
                    # 创建角色
                    role = await self.db_service.create_custom_role(
                        id=role_id,
                        user_id=DEFAULT_USER_ID,
                        name=role_data.get('name', ''),
                        title=role_data.get('title', ''),
                        input_spec=role_data.get('input_spec', ''),
                        output_spec=role_data.get('output_spec', ''),
                        custom_traits=traits,
                        system_prompt_snapshot=role_data.get('system_prompt_snapshot', ''),
                        usage_count=role_data.get('usage_count', 0)
                    )
                    
                    if role:
                        print(f"  ✅ 迁移角色: {role.name}")
                        self.stats['roles_migrated'] += 1
                    else:
                        print(f"  ❌ 迁移角色失败: {role_data.get('name', role_id)}")
                        self.stats['errors'].append(f"角色迁移失败: {role_data.get('name', role_id)}")
                        
                except Exception as e:
                    print(f"  ❌ 迁移角色出错 '{role_data.get('name', role_id)}': {e}")
                    self.stats['errors'].append(f"角色迁移出错: {role_data.get('name', role_id)} - {e}")
                    
        except Exception as e:
            print(f"❌ 读取角色文件失败: {e}")
            self.stats['errors'].append(f"读取角色文件失败: {e}")
            
        print()
    
    async def _migrate_sessions(self):
        """迁移会话和消息"""
        print("迁移会话和消息...")
        
        sessions_file = self.data_dir / "sessions.json"
        if not sessions_file.exists():
            print("⚠️  跳过会话迁移 - 文件不存在")
            return
            
        try:
            with open(sessions_file, 'r', encoding='utf-8') as f:
                sessions_data = json.load(f)
                
            for session_id, session_data in sessions_data.items():
                try:
                    # 检查会话是否已存在
                    if not self.force:
                        existing_session = await self.db_service.get_session(session_id)
                        if existing_session:
                            print(f"  ⏭️  跳过会话 '{session_data.get('title', session_id)}' - 已存在")
                            self.stats['sessions_skipped'] += 1
                            continue
                    
                    # 创建会话
                    session = await self.db_service.create_session(
                        id=session_id,
                        user_id=DEFAULT_USER_ID,
                        title=session_data.get('title', '新会话'),
                        role_id=session_data.get('role_id'),
                        model=session_data.get('model', 'deepseek-v3-1-250821'),
                        search_enabled=session_data.get('search_enabled', False),
                        memo=session_data.get('memo'),
                        memo_enabled=session_data.get('memo_enabled', False),
                        is_favorite=session_data.get('is_favorite', False),
                        is_archived=session_data.get('is_archived', False)
                    )
                    
                    if not session:
                        print(f"  ❌ 创建会话失败: {session_data.get('title', session_id)}")
                        self.stats['errors'].append(f"会话创建失败: {session_data.get('title', session_id)}")
                        continue
                        
                    print(f"  ✅ 迁移会话: {session.title}")
                    self.stats['sessions_migrated'] += 1
                    
                    # 迁移消息
                    messages = session_data.get('messages', [])
                    for i, message_data in enumerate(messages):
                        try:
                            message = await self.db_service.add_message(
                                session_id=session.id,
                                role=message_data.get('role', 'user'),
                                content=message_data.get('content', ''),
                                model=message_data.get('model'),
                                message_index=i + 1
                            )
                            
                            if message:
                                self.stats['messages_migrated'] += 1
                            else:
                                self.stats['errors'].append(f"消息创建失败: 会话 {session.title}, 消息 {i+1}")
                                
                        except Exception as e:
                            print(f"    ❌ 迁移消息出错: {e}")
                            self.stats['errors'].append(f"消息迁移出错: 会话 {session.title}, 消息 {i+1} - {e}")
                            
                except Exception as e:
                    print(f"  ❌ 迁移会话出错 '{session_data.get('title', session_id)}': {e}")
                    self.stats['errors'].append(f"会话迁移出错: {session_data.get('title', session_id)} - {e}")
                    
        except Exception as e:
            print(f"❌ 读取会话文件失败: {e}")
            self.stats['errors'].append(f"读取会话文件失败: {e}")
            
        print()
    
    def _show_migration_results(self):
        """显示迁移结果"""
        print("=== 迁移完成 ===")
        print(f"✅ 成功迁移会话: {self.stats['sessions_migrated']} 个")
        print(f"✅ 成功迁移消息: {self.stats['messages_migrated']} 条")  
        print(f"✅ 成功迁移角色: {self.stats['roles_migrated']} 个")
        
        if self.stats['sessions_skipped'] > 0:
            print(f"⏭️  跳过会话: {self.stats['sessions_skipped']} 个 (已存在)")
        if self.stats['roles_skipped'] > 0:
            print(f"⏭️  跳过角色: {self.stats['roles_skipped']} 个 (已存在)")
            
        if self.stats['errors']:
            print(f"\n❌ 遇到 {len(self.stats['errors'])} 个错误:")
            for error in self.stats['errors'][:10]:  # 只显示前10个错误
                print(f"  - {error}")
            if len(self.stats['errors']) > 10:
                print(f"  ... 还有 {len(self.stats['errors']) - 10} 个错误")
        else:
            print("\n🎉 迁移完成，没有错误！")


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='MindMate 数据库迁移工具')
    parser.add_argument('data_dir', nargs='?', default='./data', 
                       help='数据目录路径 (默认: ./data)')
    parser.add_argument('--force', action='store_true',
                       help='强制覆盖已存在的数据')
    
    args = parser.parse_args()
    
    migrator = DataMigrator(args.data_dir, args.force)
    success = await migrator.migrate()
    
    if success:
        print("\n迁移任务完成！")
        sys.exit(0)
    else:
        print("\n迁移任务失败！")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())